import os
import sys

sys.path.append(os.path.dirname(__file__))
from fastapi import FastAPI, HTTPException, Depends
from sqlalchemy.orm import Session
from db_config import SessionLocal, engine
from schemas import Text2SQLRequest, Text2SQLResponse, QueryHistory
from llm_client import LLMClient
import sqlalchemy
import re
import json
from datetime import datetime

app = FastAPI()
llm_client = LLMClient()

LOG_FILE = "text2sql_fastapi/server_debug.log"


def log_debug(msg):
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(f"[{datetime.now().isoformat()}] {msg}\n")


# 获取数据库 session
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def get_all_table_info() -> str:
    # 自动获取所有表结构，拼接为字符串
    insp = sqlalchemy.inspect(engine)
    tables = insp.get_table_names()
    info = []
    for table in tables:
        columns = insp.get_columns(table)
        col_str = ", ".join([f"{col['name']}({col['type']})" for col in columns])
        info.append(f"表: {table} 字段: {col_str}")
    return "\n".join(info)


def build_prompt(question: str, table_info: str = "") -> str:
    prompt = f"""
你是一个专业的SQL生成助手。请根据用户问题、表结构信息，生成安全的MySQL查询语句，仅限SELECT。

表结构信息：
{table_info}

【极其重要】只允许输出一条可直接执行的SQL语句，绝对禁止输出任何解释、注释、分析、思考、前后缀、代码块标记、自然语言说明，否则会导致执行失败！只输出SQL语句本身，末尾加分号。
"""
    return prompt


def extract_all_sql(text: str):
    # 提取所有 select ... ; 语句
    matches = re.findall(r"(select[\s\S]+?;)", text, re.IGNORECASE)
    sqls = [m.strip() for m in matches if m.strip().lower().startswith("select")]
    # 兜底：每行以select开头的内容
    if not sqls:
        for line in text.splitlines():
            if line.strip().lower().startswith("select"):
                sqls.append(line.strip())
    # 再兜底：只保留select到分号之间
    if not sqls:
        match = re.search(r"select[\s\S]+?;", text, re.IGNORECASE)
        if match:
            sqls.append(match.group(0).strip())
    return sqls


def log_history(history: QueryHistory):
    with open("text2sql_fastapi/query_history.log", "a", encoding="utf-8") as f:
        f.write(json.dumps(history.dict(), ensure_ascii=False) + "\n")


@app.post("/text2sql/query", response_model=Text2SQLResponse)
def text2sql_query(req: Text2SQLRequest, db: Session = Depends(get_db)):
    # 1. 获取表结构
    table_info = get_all_table_info()
    # 2. 构建 prompt
    prompt = build_prompt(req.question, table_info)
    log_debug(f"Prompt: {prompt}")
    # 3. LLM 生成 SQL
    sql_text = llm_client.generate_sql(prompt)
    log_debug(f"Model output: {sql_text}")
    sql_candidates = extract_all_sql(sql_text)
    log_debug(f"Extracted SQLs: {sql_candidates}")
    errors = []
    for sql in sql_candidates:
        if not sql.lower().strip().startswith("select"):
            continue
        # 4. 支持分页
        page = req.page or 1
        page_size = req.page_size or 10
        offset = (page - 1) * page_size
        paged_sql = sql.rstrip(';') + f" LIMIT {page_size} OFFSET {offset};"
        log_debug(f"Try SQL: {paged_sql}")
        # 5. 执行 SQL
        try:
            result_proxy = db.execute(sqlalchemy.text(paged_sql))
            result = [dict(row) for row in result_proxy]
            # 获取总数
            count_sql = f"SELECT COUNT(*) as total FROM ( {sql.rstrip(';')} ) t"
            total = db.execute(sqlalchemy.text(count_sql)).fetchone()[0]
            log_debug(f"Success SQL: {paged_sql} | Result: {result}")
            # 6. LLM 生成自然语言答案
            answer_prompt = f"根据以下SQL查询结果，用简洁中文回答用户问题：\n用户问题：{req.question}\n结果：{result}\n【只输出简洁中文答案，不要有任何解释、分析、前后缀】"
            answer = llm_client.generate_sql(answer_prompt)
            # 7. 记录历史
            history = QueryHistory(
                question=req.question,
                sql=sql,
                answer=answer,
                timestamp=datetime.now().isoformat(),
                row_count=len(result)
            )
            log_history(history)
            return Text2SQLResponse(
                sql=sql,
                result=result,
                answer=answer,
                total=total,
                page=page,
                page_size=page_size
            )
        except Exception as e:
            log_debug(f"Fail SQL: {paged_sql} | Error: {e}")
            errors.append(f"SQL: {sql}\n错误: {e}")
    log_debug(f"All SQL failed. Model output: {sql_text} | Errors: {errors}")
    raise HTTPException(status_code=500,
                        detail=f"所有SQL尝试均失败。模型输出: {sql_text}\n尝试过的SQL及错误: {' | '.join(errors)}")
